from ultralytics import YOLO
import os
import cv2
import numpy as np
import platform
import json
import torch
from text_det import TextDetector
from text_angle_cls import TextClassifier
from text_rec import TextRecognizer

print("modules_loaded")
modules_loaded = True
current_dir = os.path.dirname(os.path.abspath(__file__))
print(f"track Current directory: {current_dir}")
model_path = os.path.join(current_dir, 'weights', 'yolo11n.pt')
print(f"Model path: {model_path}")
model = YOLO(model_path)
detect_model = TextDetector()
angle_model = TextClassifier()
rec_model = TextRecognizer()
if platform.system() == 'Linux':
    device = 'cuda'
else:
    device = 'mps'
print(f"device: {device}")
initialized = True


def current_directory():
    return model_path


def process_touch(byte_array, is_text):
    global model, initialized, device, detect_model, angle_model, rec_model
    if not modules_loaded or not initialized:
        print("没创建成功")
        return json.dumps({})
    image_np = np.frombuffer(byte_array, dtype=np.uint8)
    frame = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
    if is_text:
        box_list = detect_model.detect(frame)
        texts = []
        xyxys = []
        if len(box_list) > 0:
            for point in box_list:
                point = detect_model.order_points_clockwise(point)
                textimg = detect_model.get_rotate_crop_image(frame, point.astype(np.float32))
                angle = angle_model.predict(textimg)
                if angle == '180':
                    textimg = cv2.rotate(textimg, 1)
                text = rec_model.predict_text(textimg)
                point = point.astype(int)
                x_min = np.min(point[:, 0])
                x_max = np.max(point[:, 0])
                y_min = np.min(point[:, 1])
                y_max = np.max(point[:, 1])
                new_point = [int(x_min), int(y_min), int(x_max), int(y_max)]
                xyxys.append(new_point)
                texts.append(text)
                # print("结果,", point, text)
        # print("Checking xyxys:", xyxys)
        # print("Checking texts:", texts)
        result = {
            "data": {
                "xyxys": xyxys,
                "texts": texts,
            }
        }
        # print(json.dumps(result))
        return json.dumps(result)
    else:
        results = model.track(frame, device=device, verbose=False)
        boxes = results[0].boxes
        cls = boxes.cls.int().cpu().tolist()
        track_ids = boxes.id.int().cpu().tolist()
        xyxys = boxes.xyxy.cpu().tolist()
        # print("Checking texts:", xyxys)
        result = {
            "data": {
                "track_ids": track_ids,
                "xyxys": xyxys,
                "cls": cls,
            }
        }
    return json.dumps(result)
